# -*- coding: utf-8 -*-

import numpy as np
from scipy.optimize import linear_sum_assignment
import random
import rospy
import sys
from pvesti.msg import UAVPose, UAVPosVel
from geometry_msgs.msg import Pose, Point, Twist, Vector3
from nav_msgs.msg import Odometry

#检测到的无人机个数
DETECT_UAV_NUM = 1
    
# vehicle_type = sys.argv[1]
# vehicle_num = int(sys.argv[2])

# UAV 信息类，用于存储 UAV 的时间戳、数量和位置
class UAVInfo:
    def __init__(self):
        self.timestamp = None
        self.num = 0
        self.positions = []

UAVInfo_pre = UAVInfo()  # 之前的 UAV 信息
UAVInfo_cur = UAVInfo()  # 当前的 UAV 信息
UAVInfo_odo=Point()  # 保存来自Odometry的信息

# 位置匹配类，用于匹配当前和之前的位置
class PositionMatcher:
    def __init__(self):
        self.previous_positions = []

    def match_positions(self, current_positions):
        # 如果之前的位置为空，则初始化并返回当前索引
        if len(self.previous_positions) == 0:
            self.previous_positions = current_positions
            return list(range(len(current_positions))), list(range(len(current_positions)))

        # 计算成本矩阵
        cost_matrix = np.zeros((len(self.previous_positions), len(current_positions)))
        for i in range(len(self.previous_positions)):
            for j in range(len(current_positions)):
                cost_matrix[i, j] = np.linalg.norm(np.array(self.previous_positions[i]) - np.array(current_positions[j]))

        # 使用匈牙利算法找到最小成本匹配
        row_idx, col_idx = linear_sum_assignment(cost_matrix)
        matched_indices = col_idx

        self.previous_positions = current_positions

        # 返回当前索引到之前索引的匹配
        return row_idx.tolist(), matched_indices.tolist()

# 卡尔曼滤波预测类，用于预测 UAV 的位置和速度
class KalmanFilterPredictor:
    def __init__(self, dt, process_noise=0.01, measurement_noise=5):
        self.x = None  # 初始状态（位置和速度）

        # 状态转移矩阵
        self.F = np.array([[1, 0, 0, dt, 0, 0],
                           [0, 1, 0, 0, dt, 0],
                           [0, 0, 1, 0, 0, dt],
                           [0, 0, 0, 1, 0, 0],
                           [0, 0, 0, 0, 1, 0],
                           [0, 0, 0, 0, 0, 1]], dtype=float)

        # 测量矩阵
        self.H = np.array([[1, 0, 0, 0, 0, 0],
                           [0, 1, 0, 0, 0, 0],
                           [0, 0, 1, 0, 0, 0]], dtype=float)

        # 初始状态协方差矩阵
        self.P = np.eye(6) * 1000

        # 测量噪声协方差矩阵
        self.R = np.eye(3) * measurement_noise

        # 过程噪声协方差矩阵
        self.Q = np.eye(6) * process_noise

        # 存储滤波后的位置和速度
        self.filtered_positions = []
        self.filtered_velocities = []

    def initialize_state(self, initial_position):
        # 初始化状态向量
        self.x = np.array([initial_position[0], initial_position[1], initial_position[2], 0, 0, 0], dtype=float)

    def predict(self, z, dt):
        # 更新时间间隔
        self.F[0, 3] = dt
        self.F[1, 4] = dt
        self.F[2, 5] = dt
        
        # 预测步骤
        self.x = np.dot(self.F, self.x)
        self.P = np.dot(np.dot(self.F, self.P), self.F.T) + self.Q

        # 测量更新步骤
        y = np.array(z) - np.dot(self.H, self.x)
        S = np.dot(np.dot(self.H, self.P), self.H.T) + self.R
        K = np.dot(np.dot(self.P, self.H.T), np.linalg.inv(S))
        self.x = self.x + np.dot(K, y)
        self.P = np.dot((np.eye(6) - np.dot(K, self.H)), self.P)

        # 存储结果
        self.filtered_positions.append((self.x[0], self.x[1], self.x[2]))
        self.filtered_velocities.append((self.x[3], self.x[4], self.x[5]))

# UAV 相对位置信息回调函数
def UAV_pose_callback(data):
    global UAVInfo_pre, UAVInfo_cur, UAVInfo_odo
    
    # 复制 UAVInfo_cur 的数据到 UAVInfo_pre
    UAVInfo_pre = UAVInfo()
    UAVInfo_pre.timestamp = UAVInfo_cur.timestamp
    UAVInfo_pre.num = UAVInfo_cur.num
    UAVInfo_pre.positions = UAVInfo_cur.positions
    
    # 更新 UAVInfo_cur 的数据
    UAVInfo_cur.timestamp = data.header.stamp
    UAVInfo_cur.num = data.uav_count
    
    #没有订阅到vins时，UAVInfo_odo直接是（0，0，0）
    UAVInfo_cur.positions = [
    (
        pose.position.x + UAVInfo_odo.x, 
        pose.position.y + UAVInfo_odo.y, 
        pose.position.z + UAVInfo_odo.z
    ) 
    for i, pose in enumerate(data.uav_positions)
]

# UAV 绝对位置信息回调函数
def vins_callback(data):
    global UAVInfo_odo
    
    UAVInfo_odo=Point()
    UAVInfo_odo=data.pose.pose.position


# 主函数，用于测试位置匹配和卡尔曼滤波
def main():
    matcher = PositionMatcher()

    dt = 1.0
    kalman_predictors = {}
    kalman_index_map = {}  # 当前索引到初始索引的映射

    # 测试数据
    m = 5  # 时间步数
    n = 3  # 物体数量
    positions = [[(1,1,2), (3,3,2), (6,1,2)],
                 [(3.2, 3.2, 2), (1.7, 1.8, 2), (6.1, 1.8, 2)],
                 [(6.4, 1.7, 2), (1.2, 1.2, 2), (3.9, 3.9, 2)],
                 [(1.5, 1.5, 2), (3.3, 3.4, 2), (6.7, 1.2, 2)],
                 [(1.3, 1.7, 2), (6.8, 1.6, 2), (3.2, 3.9, 2)]]

    initial_indices = list(range(n))
    for i in range(m):
        positions_cur = positions[i]

        # 匹配位置
        row_index, matched_index = matcher.match_positions(positions_cur)

        if i == 0:
            # 初始化卡尔曼滤波器
            for idx in initial_indices:
                kalman_predictors[idx] = KalmanFilterPredictor(dt)
                kalman_predictors[idx].initialize_state(positions_cur[idx])
                kalman_index_map[idx] = idx
        else:
            # 更新卡尔曼滤波器索引映射
            #！！关键转换！！
            kalman_index_map_new = {}
            for j, idx in zip(row_index, matched_index):
                #当前匹配中的last index对应的初始坐标
                initial_index = kalman_index_map[j]
                kalman_index_map_new[idx] = initial_index
            kalman_index_map = kalman_index_map_new

        for value, key in kalman_index_map.items():
            kalman_predictors[key].predict(positions_cur[value], dt)

    # 打印滤波后的位置和速度
    for i in kalman_predictors:
        filtered_positions = [(float(x), float(y), float(z)) for (x, y, z) in kalman_predictors[i].filtered_positions]
        filtered_velocities = [(float(vx), float(vy), float(vz)) for (vx, vy, vz) in kalman_predictors[i].filtered_velocities]
        print("Object {} filtered positions: {}\n".format(i, filtered_positions))
        print("Object {} filtered velocities: {}\n".format(i, filtered_velocities))


# 主函数2，用于 ROS 节点，接收 UAV 位置信息并发布位置信息和速度
def main2():
    rospy.init_node('pos_vel_transfer')
    rospy.Subscriber('/vins_estimator/odometry', Odometry,vins_callback, queue_size=1)#从vins中得到中心无人机的绝对位置信息
    # rospy.Subscriber(vehicle_type+'_'+str(i)+'/vins_estimator/odometry', Odometry,vins_callback, queue_size=1)#从vins中得到中心无人机的绝对位置信息
    rospy.Subscriber('/uav_positions', UAVPose, UAV_pose_callback, queue_size=1)#从雷达发布节点处得到无人机数量、相对位置信息uav_count、geometry_msgs/Pose
    pub = rospy.Publisher('/uav_pos_vel', UAVPosVel, queue_size=10)#给处理避障的节点发布无人机数量、位置、速度信息uav_count、geometry_msgs/Pose、geometry_msgs/Twist
    

    matcher = PositionMatcher()
    kalman_predictors = {}
    kalman_index_map = {}
    #新增，单独的卡尔曼滤波器用于当前无人机
    current_uav_kalman_predictor = KalmanFilterPredictor(dt=1.0)
    
    global UAVInfo_pre, UAVInfo_cur, UAVInfo_odo
    
    first_run = True
    rate = rospy.Rate(1)
    
    
    # 卡尔曼的估计频率就是dt，但是是预测速度+
    while not rospy.is_shutdown():
        if UAVInfo_cur.timestamp is not None and UAVInfo_pre.timestamp is not None:
            dt = UAVInfo_cur.timestamp.to_sec() - UAVInfo_pre.timestamp.to_sec()
            positions_cur = UAVInfo_cur.positions
            
            if first_run:
                for idx in range(DETECT_UAV_NUM):
                    kalman_predictors[idx] = KalmanFilterPredictor(dt)
                    kalman_predictors[idx].initialize_state(positions_cur[idx])
                    kalman_index_map[idx] = idx
                # 当前uav卡尔曼滤波器-初始化    
                current_uav_kalman_predictor.initialize_state([UAVInfo_odo.x, UAVInfo_odo.y, UAVInfo_odo.z])
                first_run = False
            else:
                row_index, matched_index = matcher.match_positions(positions_cur)
                kalman_index_map_new = {}
                for j, idx in zip(row_index, matched_index):
                    initial_index = kalman_index_map[j]
                    kalman_index_map_new[idx] = initial_index
                kalman_index_map = kalman_index_map_new
            
            for value, key in kalman_index_map.items():
                kalman_predictors[key].predict(positions_cur[value], dt)
            
            # 当前uav卡尔曼滤波器-预测
            current_position = [UAVInfo_odo.x, UAVInfo_odo.y, UAVInfo_odo.z]
            current_uav_kalman_predictor.predict(current_position, dt)
               
            msg = UAVPosVel()
            msg.header.stamp = UAVInfo_cur.timestamp
            msg.uav_count = DETECT_UAV_NUM
            
            # 增加当前无人机的位置和速度
            pose_cur = Pose()
            pose_cur.position = Point(*current_uav_kalman_predictor.filtered_positions[-1])
            twist_cur = Twist()
            twist_cur.linear = Vector3(*current_uav_kalman_predictor.filtered_velocities[-1])
            msg.uav_positions.append(pose_cur)
            msg.uav_velocities.append(twist_cur)
            
            pose_cur_info = [pose_cur.position.x, pose_cur.position.y, pose_cur.position.z]
            vel_cur_info = [twist_cur.linear.x, twist_cur.linear.y, twist_cur.linear.z]
            rospy.loginfo("current UAV position': {}".format(pose_cur_info))
            rospy.loginfo("current UAV velocity': {}".format(vel_cur_info))
            
            
            for i in range(DETECT_UAV_NUM):
                pose = Pose()
                pose.position = Point(*kalman_predictors[i].filtered_positions[-1])
                twist = Twist()
                twist.linear = Vector3(*kalman_predictors[i].filtered_velocities[-1])
                msg.uav_positions.append(pose)
                msg.uav_velocities.append(twist)
                
                pose_info = [pose.position.x, pose.position.y, pose.position.z]
                vel_info = [twist.linear.x, twist.linear.y, twist.linear.z]
                
                
                rospy.loginfo("UAV {} position': {}".format(i, pose_info))
                rospy.loginfo("UAV {} velocity': {}".format(i, vel_info))
                
                
            pub.publish(msg)
            
        rate.sleep()

if __name__ == "__main__":
    main()